"""
Created by joe on 2019/12/13

多训练集合并
"""
import os
import shutil
import uuid

img_subs = ["jpg", "jpeg", "png"]


class TrainDirMerge(object):
    """多训练集合并"""

    def __init__(self, *dirs):
        """"""
        self.classes = []
        self.positions = {}

        for dir_path in dirs:
            if dir_path:
                self.cache(dir_path)

    def cache(self, dir_path):
        """

        :param dir_path:
        :return:
        """
        child_classes = []
        with open("{0}/classes.txt".format(dir_path), "r") as fi:
            for line in fi.readlines():
                line = line.strip()
                child_classes.append(line)

        filenames = os.listdir(dir_path)
        for filename in filenames:
            name_arr = filename.split(".")

            if name_arr[1] not in img_subs:
                continue

            img_path = "{0}/{1}".format(dir_path, filename)

            with open("{0}/{1}.txt".format(dir_path, name_arr[0]), "r") as fi:
                for line in fi.readlines():
                    if img_path not in self.positions:
                        self.positions[img_path] = []

                    line = line.strip()

                    if not line:
                        continue

                    index = line.index(" ")
                    class_num = int(line[:index])
                    position = line[index:].strip()
                    class_name = child_classes[class_num]
                    if class_name not in self.classes:
                        self.classes.append(class_name)
                    self.positions[img_path].append((self.classes.index(class_name), position))

    def merge_to(self, target_path):
        """"""
        if not os.path.exists(target_path):
            os.mkdir(target_path)

        with open("{0}/classes.txt".format(target_path), "w") as fi:
            for line in self.classes:
                fi.write(line + "\n")

        for org_path in self.positions:
            index2 = org_path.rindex(".")
            suffix = org_path[index2 + 1:]
            new_name = str(uuid.uuid4())

            new_path = "{0}/{1}.{2}".format(target_path, new_name, suffix)
            shutil.copy(org_path, new_path)

            with open("{0}/{1}.txt".format(target_path, new_name), "w") as fi:
                position = self.positions[org_path]
                for line in position:
                    fi.write("{0} {1}".format(line[0], line[1]) + "\n")


if __name__ == "__main__":
    """main"""
    obj = TrainDirMerge("/home/joe/Desktop/captcha_01", "/home/joe/Desktop/captcha_02", "/home/joe/Desktop/captcha_04",
                        "/home/joe/Desktop/captcha_05", "/home/joe/Desktop/captcha_06")
    obj.merge_to("/home/joe/Desktop/captcha_merge")
